import logging

import akshare as ak
import torch
from torch import nn

from StockDataSource import StockDataSource
from StockDataTrain import StockLSTM

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    filename='StockDataExecute.log', )

device = torch.device("cpu")

torch.set_num_threads(4)


def train_step(model, criterion, optimizer, x, y):
    y_pred = model(x)
    loss = criterion(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()


def train_model(model, criterion, optimizer, train_loader):
    model.train()
    train_loss = 0
    for x, y in train_loader:
        train_loss += train_step(model, criterion, optimizer, x, y)
    return train_loss / len(train_loader)


def test_model(model, criterion, test_loader):
    model.eval()
    epoch_loss = 0
    with torch.no_grad():
        for x, y in test_loader:
            test_outputs = model(x).squeeze(1)
            test_loss = criterion(test_outputs, y)
            epoch_loss += test_loss.item()
    return epoch_loss / len(test_loader)


if __name__ == '__main__':
    model = StockLSTM(16, 3, 1, 6)
    model = model.to(device)
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    stock_list = ak.stock_zh_a_spot_em()
    new_list1 = stock_list[stock_list["代码"].str.contains("^60")]
    new_list2 = stock_list[stock_list["代码"].str.contains("^00")]
    new_list3 = stock_list[stock_list["代码"].str.contains("^30")]

    code_list = []
    for code in new_list1["代码"].values:
        code_list.append(code)

    for code in new_list2["代码"].values:
        code_list.append(code)

    for code in new_list3["代码"].values:
        code_list.append(code)

    for stock_code in code_list:
        try:
            stockDataSource = StockDataSource(stock_code=stock_code, days=3)
            train_data_loader, test_data_loader = stockDataSource.getTrainData()

            epochs = 1000
            for epoch in range(epochs):
                train_loss = train_model(model, criterion, optimizer, train_data_loader)
                eval_loss = test_model(model, criterion, test_data_loader)
                logging.info(
                    f'StockCode: {stock_code}, Epoch: {epoch + 1}, Train Loss: {train_loss:.6f}, Eval Loss: {eval_loss:.6f}')
        except Exception as e:
            logging.error("occur exception: ", e)
